from PIL import Image
import os
import xml.etree.ElementTree as ET
from data_check import check_datasets


# """
#     @author: xy
#     @description: 图片格式 图片标签处理
# """
#
# data_root = "MaskPascalVOC"
# image_root = "JPEGImages"
# annotations_root = "Annotations"
#
# image_format_wrong = []
#
# # 检查数据格式是否为png或jpg、检查图片文件与标签文件数量是否对应、统计图片和标签数量
# images = os.listdir(os.path.join(data_root, image_root))
# for image in images:
#     image_path = os.path.join(os.path.join(data_root, image_root), image)
#     img = Image.open(image_path)
#     if img.format not in ["JPEG", "JPG", "jpg", "jpeg"] and img.format in ['png', "PNG"]:
#         # 记录是png但不是jpg的图片名称
#         image_format_wrong.append(image)
#     elif img.format in ["JPEG", "JPG", "jpeg", "jpg"]:
#         continue
#     else:
#         # 删除既不是png也不是jpg格式的图片和对应的标签
#         try:
#             os.remove(os.path.join(os.path.join(data_root, image_root), image))
#             os.remove(os.path.join(os.path.join(data_root, annotations_root), img.split('.')[0] + ".xml"))
#         except:
#             continue
#
# # 将是png但不是jpg的图片的格式转为jpg格式并删除原png格式的图片及对应的标签
# for image in image_format_wrong:
#     image_path = os.path.join(os.path.join(data_root, image_root), image)
#     img = Image.open(image_path)
#     im = img.convert('RGB')
#     im.save(os.path.join(os.path.join(data_root, image_root), im.split('.')[0] + ".jpg"), quality=100)
#     os.remove(os.path.join(os.path.join(data_root, image_root), image))
#     print("process png image successful")
#

# 修改png格式的图片为jpg格式的图片
def pngs2jpgs(root):
    input_folder = "JPEGImages"
    output_folder = "JPEGImages"
    annotations_folder = "Annotations"
    png_path = os.path.join(root, output_folder)
    imgs = os.listdir(png_path)
    jpg_path = os.path.join(root, output_folder)
    for img in imgs:
        img_path = os.path.join(png_path, img)
        image = Image.open(img_path)
        if image.format not in ["JPEG", "JPG", "jpg", "jpeg"] and image.format in ['png', "PNG"]:
            print("Image '{}' format not JPEG/JPG".format(img_path))
            print(image.format)
            im = image.convert('RGB')
            jpg_img_name = img.split('.')[0] + ".jpg"
            im.save(os.path.join(jpg_path, jpg_img_name), quality=100)
            # os.remove(os.path.join(png_path, img))
            # os.remove(os.path.join(os.path.join(root, annotations_folder), img.split('.')[0] + ".xml"))
            print("successful png2jpg")
        elif image.format in ["JPEG", "JPG", "jpeg", "jpg"]:
            continue
        else:
            # os.remove(os.path.join(png_path, img))
            # os.remove(os.path.join(os.path.join(root, annotations_folder), img.split('.')[0] + ".xml"))
            print("remove img not png and jpg")


# 将图片目录和标签目录中的文件重命名 并修改xml文件中的name标签
def name2seq(root):
    xml_path = os.path.join(root, "Annotations")
    jpg_path = os.path.join(root, "JPEGImages")
    xmls = os.listdir(xml_path)
    jpgs = os.listdir(jpg_path)
    if len(xmls) == len(jpgs):
        nums = len(xmls)
        print(nums)
    else:
        nums = min(len(xmls), len(jpgs))
        print("标签与图片数量不相等！")
        # check_datasets(root)
    name_seq = list(range(nums))
    print(len(name_seq))
    for i in range(len(name_seq)):
        os.rename(os.path.join(xml_path, xmls[i]), os.path.join(xml_path, str("CFM") + str(i) + ".xml"))
        os.rename(os.path.join(jpg_path, jpgs[i]), os.path.join(jpg_path, str("CFM") + str(i) + ".jpg"))
    # xmls2jpgxmls(root)


# 修改xml文件中的name为xml文件名 + .jpg
def xmls2jpgxmls(root):
    xml_path = os.path.join(root, "Annotations")
    xmls = os.listdir(xml_path)
    print(xmls)
    # 打开xml文档
    for xml in xmls:
        doc = ET.parse(os.path.join(xml_path, xml))
        root = doc.getroot()
        sub = root.find('filename')  # 找到filename标签，
        sub.text = xml.split(".")[0] + ".jpg"  # 修改标签内容
        doc.write(os.path.join(xml_path, xml))  # 保存修改
    for xml in xmls:
        doc = ET.parse(os.path.join(xml_path, xml))
        root = doc.getroot()
        for child in root:
            if child.tag == 'path' or child.tag == 'folder':
                root.remove(child)
        doc.write(os.path.join(xml_path, xml))  # 保存修改


def xmls_process(root):
    xml_path = os.path.join(root, "Annotations")
    jpg_path = os.path.join(root, "JPEGImages")
    xmls = os.listdir(xml_path)
    # print(xmls)
    # 打开xml文档
    nametext_list = []
    for xml in xmls:
        doc = ET.parse(os.path.join(xml_path, xml))
        root = doc.getroot()
        tag_list = []
        for child in root:
            tag_list.append(child.tag)
            if child.tag == 'object':
                for child_child in child:
                    if child_child.tag == 'bndbox':
                        xmin = child_child.find('xmin').text
                        ymin = child_child.find('ymin').text
                        xmax = child_child.find('xmax').text
                        ymax = child_child.find('ymax').text
                        # print(xmax, xmin)
                        if abs(float(xmax) - float(xmin)) < 3 or abs(float(ymax) - float(ymin)) < 3 \
                                or float(xmax) <= float(xmin) or float(ymax) <= float(ymin):
                            # 删掉标记框脏数据
                            try:
                                os.remove(os.path.join(xml_path, xml))
                                jpg_name = xml.split(".")[0] + ".jpg"
                                os.remove(os.path.join(jpg_path, jpg_name))
                                print("remove bndbox w/h<=0 file")
                                continue
                            except FileNotFoundError as e:
                                print(os.path.join(xml_path, xml) + "FileNotFoundError: ")
                                continue
        if 'object' not in tag_list:
            os.remove(os.path.join(xml_path, xml))
            jpg_name = xml.split(".")[0] + ".jpg"
            os.remove(os.path.join(jpg_path, jpg_name))
            print("remove xml without object")


# 统一所有标签文件中的标签
def xmls2same(root):
    xml_path = os.path.join(root, "Annotations")
    jpg_path = os.path.join(root, "JPEGImages")
    xmls = os.listdir(xml_path)
    # print(xmls)
    # 打开xml文档
    nametext_list = []
    nums = 0
    for xml in xmls:
        doc = ET.parse(os.path.join(xml_path, xml))
        root = doc.getroot()
        two_tag = 0
        for child in root:
            if child.tag == 'object':
                sub = child.find('name')
                if sub.text == 'face_mask':
                    # 直接舍弃 删掉同时标记人脸和口罩的数据集
                    two_tag = 1
            if child.tag == "object":
                sub = child.find('name')  # 找到object->name标签
                for child_child in child:
                    if child_child.text == "none":
                        child_child.text = "none_mask"
                    if child_child.text == "mask_weared_incorrect":
                        child_child.text = "poor_mask"
                    if child_child.text == "without_mask":
                        child_child.text = "none_mask"
                    if child_child.text == "face_mask":
                        child_child.text = "with_mask"
                    if child_child.text == "face":
                        child_child.text = "none_mask"
                    if child_child.text == "poor":
                        child_child.text = "poor_mask"
                    if child_child.text == "nomask":
                        child_child.text = "none_mask"
                    if child_child.text == "mask":
                        child_child.text = "with_mask"
                nametext_list.append(sub.text)
        if two_tag == 1:
            nums = nums + 1
            # 删掉同时标记人脸和口罩的数据集
            os.remove(os.path.join(xml_path, xml))
            jpg_name = xml.split(".")[0] + ".jpg"
            os.remove(os.path.join(jpg_path, jpg_name))
            print("remove two_tag file")
        else:
            doc.write(os.path.join(xml_path, xml))  # 保存修改
    print(set(nametext_list))
    print(len(nametext_list))
    print(nums)


def delete_addition(root):
    xml_path = os.path.join(root, "Annotations")
    jpg_path = os.path.join(root, "JPEGImages")
    xmls = os.listdir(xml_path)
    # print(xmls)
    # 打开xml文档
    only_one_with_mask = 0
    only_one_none_mask = 0
    only_one_poor_mask = 0
    for xml in xmls:
        with_mask_list = []
        poor_mask_list = []
        none_mask_list = []
        doc = ET.parse(os.path.join(xml_path, xml))
        root = doc.getroot()
        for child in root:
            if child.tag == 'object':
                sub = child.find('name')
                for child_child in child:
                    if child_child.text == "with_mask":
                        with_mask_list.append(child_child.text)
                        print(child_child.text)
                    elif child_child.text == "poor_mask":
                        poor_mask_list.append(child_child.text)
                        print(child_child.text)
                    elif child_child.text == "none_mask":
                        none_mask_list.append(child_child.text)
                        print(child_child.text)
        if len(with_mask_list) == 1 and len(poor_mask_list) == 0 and len(none_mask_list) == 0:
            only_one_with_mask = only_one_with_mask + 1
            # 删掉同时标记人脸和口罩的数据集
            # os.remove(os.path.join(xml_path, xml))
            # jpg_name = xml.split(".")[0] + ".jpg"
            # os.remove(os.path.join(jpg_path, jpg_name))
        if len(with_mask_list) == 0 and len(poor_mask_list) == 1 and len(none_mask_list) == 0:
            only_one_poor_mask = only_one_poor_mask + 1
        if len(with_mask_list) == 0 and len(poor_mask_list) == 0 and len(none_mask_list) == 1:
            only_one_none_mask = only_one_none_mask + 1

        # elif len(none_mask_list) == 1 and len(poor_mask_list) == 0 and len(with_mask_list) == 0:
        #     only_one_none_mask = only_one_none_mask + 1
        #     # 删掉同时标记人脸和口罩的数据集
        #     os.remove(os.path.join(xml_path, xml))
        #     jpg_name = xml.split(".")[0] + ".jpg"
        #     os.remove(os.path.join(jpg_path, jpg_name))
        # elif len(poor_mask_list) == 1 and len(none_mask_list) == 0 and len(with_mask_list) == 0:
        #     only_one_poor_mask = only_one_poor_mask + 1
        #     # 删掉同时标记人脸和口罩的数据集
        #     os.remove(os.path.join(xml_path, xml))
        #     jpg_name = xml.split(".")[0] + ".jpg"
        #     os.remove(os.path.join(jpg_path, jpg_name))
    print(only_one_with_mask)
    print(only_one_poor_mask)
    print(only_one_none_mask)


def find_poor_mask(root):
    xml_path = os.path.join(root, "Annotations")
    jpg_path = os.path.join(root, "JPEGImages")
    xmls = os.listdir(xml_path)
    jpgs = os.listdir(jpg_path)
    with_mask = 0
    none_mask = 0
    poor_mask = 0
    for xml in xmls:
        with_mask_list = []
        poor_mask_list = []
        none_mask_list = []
        doc = ET.parse(os.path.join(xml_path, xml))
        root = doc.getroot()
        for child in root:
            if child.tag == 'object':
                sub = child.find('name')
                for child_child in child:
                    if child_child.text == "with_mask":
                        with_mask_list.append(child_child.text)
                        with_mask = with_mask + 1
                    elif child_child.text == "poor_mask":
                        poor_mask_list.append(child_child.text)
                        poor_mask = poor_mask + 1
                    elif child_child.text == "none_mask":
                        none_mask_list.append(child_child.text)
                        none_mask = none_mask + 1
    print(with_mask)
    print(poor_mask)
    print(none_mask)
    times = max(none_mask, with_mask)
    for xml in xmls:
        if (none_mask > 500 or with_mask > 500) and times > 0:
            times = times - 1
            with_mask_list = []
            poor_mask_list = []
            none_mask_list = []
            doc = ET.parse(os.path.join(xml_path, xml))
            root = doc.getroot()
            for child in root:
                if child.tag == 'object':
                    sub = child.find('name')
                    for child_child in child:
                        if child_child.text == "with_mask":
                            with_mask_list.append(child_child.text)
                        elif child_child.text == "poor_mask":
                            poor_mask_list.append(child_child.text)
                        elif child_child.text == "none_mask":
                            none_mask_list.append(child_child.text)
                            print("get")
            # if len(poor_mask_list) == 0 and len(none_mask_list) > 0:
            #     none_mask = none_mask - len(none_mask_list)
            #     jpg_name = xml.split(".")[0] + ".jpg"
            #     os.remove(os.path.join(xml_path, xml))
            #     os.remove(os.path.join(jpg_path, jpg_name))
            #     print("remove")
            if len(with_mask_list) > 0:
                with_mask = with_mask - len(with_mask_list)
                jpg_name = xml.split(".")[0] + ".jpg"
                os.remove(os.path.join(xml_path, xml))
                os.remove(os.path.join(jpg_path, jpg_name))
                print("remove")

        # if len(poor_mask_list) == 0 and len(with_mask_list) > 1 and len(none_mask_list) > 1:
        #     only_one_with_mask = only_one_with_mask + 1
        #     # 删掉同时标记人脸和口罩的数据集
        #     os.remove(os.path.join(xml_path, xml))
        #     jpg_name = xml.split(".")[0] + ".jpg"
        #     os.remove(os.path.join(jpg_path, jpg_name))
        # if len(poor_mask_list) == 0 and len(with_mask_list) > 1 and len(none_mask_list) > 1:
        #     only_one_with_mask = only_one_with_mask + 1
        #     # 删掉同时标记人脸和口罩的数据集
        #     os.remove(os.path.join(xml_path, xml))
        #     jpg_name = xml.split(".")[0] + ".jpg"
        #     os.remove(os.path.join(jpg_path, jpg_name))

def find_other_folder(root):
    xml_path = os.path.join(root, "Annotations")
    jpg_path = os.path.join(root, "JPEGImages")
    xmls = os.listdir(xml_path)
    jpgs = os.listdir(jpg_path)
    file_name_list = []
    xml_name_list = []
    for xml in xmls:
        with_mask_list = []
        poor_mask_list = []
        none_mask_list = []
        doc = ET.parse(os.path.join(xml_path, xml))
        root = doc.getroot()
        for child in root:
            if child.tag == 'object':
                sub = child.find('name')
                for child_child in child:
                    if child_child.text == "with_mask":
                        with_mask_list.append(child_child.text)
                    elif child_child.text == "poor_mask":
                        poor_mask_list.append(child_child.text)
                    elif child_child.text == "none_mask":
                        none_mask_list.append(child_child.text)
        if len(poor_mask_list) > 2:
            file_name_list.append(os.path.join(jpg_path, xml))
            xml_name_list.append(os.path.join(xml_path, xml.split('.')[0]+'.jpg'))
    print(len(file_name_list))
    print(xml_name_list)


if __name__ == '__main__':
    datasets_folder_root = 'F:/face-mask-detection/MaskDatasets-Light'
    # pngs2jpgs(datasets_folder_root)
    # name2seq(datasets_folder_root)
    # xmls2jpgxmls(datasets_folder_root)
    xmls_process(datasets_folder_root)
    # xmls2same(datasets_folder_root)
    # check_datasets(datasets_folder_root)
    # delete_addition(datasets_folder_root)
    # find_poor_mask(datasets_folder_root)
    # find_other_folder(r"F:\face-mask-detection\ORIimages")
